package utils

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"github.com/zeebo/blake3"
	"io"
	"os"
	"sync"
	"time"
)

const (
	// ChunkSize 定义块的大小为32MB
	ChunkSize int = 32 * 1024 * 1024 // 32MB
	// MaxMemory 定义最大内存使用量为2GB
	MaxMemory int64 = 2 * 1024 * 1024 * 1024 // 2GB
)

type ChunkResult struct {
	Index int
	Hash  []byte
}

// ComputeFileBlake3Hash 计算指定文件的BLAKE3哈希值
func computeFileBlake3Hash(filePath string) ([]byte, int64, error) {
	// 创建一个带取消功能的上下文
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel() // 确保在函数结束时取消上下文

	file, err := os.Open(filePath)
	if err != nil {
		return nil, 0, fmt.Errorf("无法打开文件: %w", err)
	}
	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	info, err := file.Stat()
	if err != nil {
		return nil, 0, fmt.Errorf("无法获取文件信息: %w", err)
	}

	totalSize := info.Size()
	totalChunks := int((totalSize + int64(ChunkSize) - 1) / int64(ChunkSize))
	//fmt.Printf("文件总大小: %d bytes, 分块数量: %d\n", totalSize, totalChunks)

	// 创建信号量，用于控制最大内存使用量
	semaphore := make(chan int64, MaxMemory/int64(ChunkSize))
	// 创建作业和结果通道
	jobs := make(chan struct {
		Index int
		Data  []byte
	})
	results := make(chan ChunkResult, 64)

	// 错误通道，用于传递错误
	errCh := make(chan error, 1)

	var wg sync.WaitGroup

	// 计算工人数量，可以根据CPU核心数动态调整
	workerCount := int(MaxMemory/int64(ChunkSize)) * 2
	for w := 0; w < workerCount; w++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				case job, ok := <-jobs:
					if !ok {
						return
					}
					// 处理哈希计算
					hasher := blake3.New()
					_, err := hasher.Write(job.Data)
					if err != nil {
						select {
						case errCh <- fmt.Errorf("worker %d: 哈希计算错误: %w", workerID, err):
							cancel()
						default:
						}
						return
					}
					hash := hasher.Sum(nil)
					select {
					case results <- ChunkResult{Index: job.Index, Hash: hash}:
					case <-ctx.Done():
						return
					}
					// 处理完毕，释放信号量
					select {
					case <-semaphore:
					case <-ctx.Done():
						return
					}
				}
			}
		}(w)
	}

	// 记录开始时间
	startTime := time.Now()

	// 读取文件并发送到jobs通道
	go func() {
		defer func() {
			close(jobs)
		}()
		reader := bufio.NewReader(file)
		for i := 0; i < totalChunks; i++ {
			select {
			case <-ctx.Done():
				return
			default:
			}

			bufferSize := ChunkSize
			if i == totalChunks-1 {
				remaining := totalSize - int64(i)*int64(ChunkSize)
				bufferSize = int(remaining)
			}
			buffer := make([]byte, bufferSize)
			_, err := io.ReadFull(reader, buffer)
			if err != nil {
				if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) {
					// 最后一块可能不足块大小，继续处理
				} else {
					select {
					case errCh <- fmt.Errorf("读取文件错误: %w", err):
						cancel()
					default:
					}
					return
				}
			}
			// 获取信号量，控制内存使用
			select {
			case semaphore <- int64(bufferSize):
			case <-ctx.Done():
				return
			}

			select {
			case jobs <- struct {
				Index int
				Data  []byte
			}{Index: i, Data: buffer}:
			case <-ctx.Done():
				return
			}
		}
	}()

	// 收集所有哈希结果
	hashes := make([][]byte, totalChunks)
	done := make(chan struct{})
	go func() {
		defer close(done)
		for res := range results {
			hashes[res.Index] = res.Hash
		}
	}()

	// 等待所有工作完成
	go func() {
		wg.Wait()
		close(results)
	}()

	// 监控错误或完成信号
	select {
	case err := <-errCh:
		return nil, 0, err
	case <-done:
		// 所有结果已收集，无错误发生
	}

	// 合并所有块的哈希来生成最终的哈希
	finalHasher := blake3.New()
	for _, hash := range hashes {
		_, err := finalHasher.Write(hash)
		if err != nil {
			return nil, 0, fmt.Errorf("最终哈希计算错误: %w", err)
		}
	}
	finalHash := finalHasher.Sum(nil)

	// 记录结束时间并计算耗时
	duration := time.Since(startTime)
	nanoSeconds := duration.Nanoseconds()
	//fmt.Printf("总耗时: %v\n", duration)

	return finalHash, nanoSeconds, nil
}

func ComputeFileHashString(filePath string) (string, int64, error) {
	hash, nanoSeconds, err := computeFileBlake3Hash(filePath)
	if err != nil {
		return "", 0, err
	}
	return fmt.Sprintf("%x", hash), nanoSeconds, nil
}

func ComputeFileHashBytes(filePath string) ([]byte, int64, error) {
	return computeFileBlake3Hash(filePath)
}
